from modelscope import AutoModelForCausalLM, AutoTokenizer, AutoModel
from typing import Optional
import torch
import os
import pickle
import torch.nn.functional as F
import chromadb
import networkx as nx
import sentence_transformers

db_path="./storege"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


os.environ['MODELSCOPE_CACHE'] = 'D:/cache/modelscope'
os.environ['HF_HOME'] = 'D:/cache/huggingface'

# 确保目录存在
os.makedirs('D:/cache/modelscope', exist_ok=True)
os.makedirs('D:/cache/huggingface', exist_ok=True)
class model:
    def __init__(self,storage_path):
        """
        初始化 GraphRAG 引擎
        修改点：
        1. 使用 ModelScope 的 embedding 模型替代 sentence-transformers
        2. 整合 mean pooling 方法
        3. 优化设备管理和内存使用

        Args:
            vector_db: ChromaDB 集合对象
            knowledge_graph: networkx 图对象
        """
        #加载知识图谱。
        self.persist_directory = storage_path
        kg_path = os.path.join(self.persist_directory, "knowledge_graph.pkl")
        if os.path.exists(kg_path):
            with open(kg_path, "rb") as f:
                self.knowledge_graph = pickle.load(f)
        #加载向量数据库，用来匹配；
            # 加载向量数据库
        self.vector_db_client = chromadb.PersistentClient(path=db_path)
        self.vector_db = self.vector_db_client.get_collection("chinese_knowledge_base")
        # 初始化 embedding 模型
        self.embedding_model = AutoModel.from_pretrained(
            "Qwen/Qwen3-Embedding-0.6B"
        )
        self.embedding_model.to(device)
        self.embedding_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-Embedding-0.6B")

        # 初始化 LLM
        self.llm = AutoModelForCausalLM.from_pretrained(
            "Qwen/Qwen3-8B",
            torch_dtype="auto",
            device_map="auto",
            load_in_4bit=True
        )
        self.llm_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")

    def _mean_pooling(self, model_output, attention_mask):
        """
        Mean pooling - 考虑 attention mask 进行平均
        修改自您提供的示例代码
        """
        token_embeddings = model_output.last_hidden_state
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

    def _get_query_embedding(self, query: str):
        """
        生成查询的 embedding 向量
        使用 ModelScope 模型 + mean pooling
        """
        # 获取模型所在的设备
        device = next(self.embedding_model.parameters()).device

        # 将输入移到相同设备
        inputs = self.embedding_tokenizer(query, return_tensors="pt", padding=True, truncation=True).to(device)

        with torch.no_grad():
            model_output = self.embedding_model(**inputs)
            embeddings = self._mean_pooling(model_output, inputs['attention_mask'])

        return embeddings[0].cpu().numpy()

    def match_kg_context(self, query: str, top_k: int = 3) -> str:
        """
        简化版相关子图信息检索
        """
        # 根据用户提问查找相关片段
        query_embedding = self._get_query_embedding(query)
        #根据问题的embedding向量在数据库里查询相似的向量，返回对应的数据库块，每一个数据库块包含存储的实体关系图或者原始文档，可以用这些字符串里的实体构建知识图谱
        results = self.vector_db.query(
            query_embeddings=[query_embedding],
            n_results=top_k
        )
        print(f"查询结果：\n {results}")
        # 遍历知识图谱，将在查找到的片段中出现的实体的所有关系列出作为模型的context;
        subgraph_info = []
        for doc in results["documents"][0]:
            # 简化实体匹配逻辑
            for entity in self.knowledge_graph.nodes:
                if entity in doc:
                    # 获取该实体相关的所有关系
                    for neighbor in self.knowledge_graph.neighbors(entity):
                        relation = self.knowledge_graph.edges[entity, neighbor].get("relation", "related_to")
                        subgraph_info.append(f"{entity} -[{relation}]-> {neighbor}")

        return "\n".join(subgraph_info) if subgraph_info else "未找到相关知识"

    def query(self, prompt: str, max_new_tokens: int = 30000) -> str:
        """
        执行 GraphRAG 查询
        修改点：
        1. 增强 prompt 模板
        2. 优化生成参数
        """
        # 1. 检索知识子图
        graph_context = self.match_kg_context(prompt)

        # 2. 构造增强 Prompt
        augmented_prompt = f"""你会学习逻辑实体关系图来学习新的编程语言，你不会c/c++、java、python、go、等常见语言，你只会根据知识图谱的语言解决问题：

[用户提问所用的编程语言的控制逻辑实体关系图]
{graph_context}

[用户问题]
{prompt}

[回答要求]
- 只使用提供的知识，根据提供的知识图谱学习这门语言，回答用户问题
- 保持简洁准确
- 避免编造信息"""

        # 3. 调用本地模型生成
        messages = [
            {"role": "system", "content": "你是一个编程大师，在回答每一个问题前你都会学习content里的知识，content里面包含着一个编程语言的控制逻辑实体关系图，只能根据提供的事实回答问题。"},
            {"role": "user", "content": augmented_prompt}
        ]
        print(graph_context)
        text = self.llm_tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )

        inputs = self.llm_tokenizer([text], return_tensors="pt").to(self.llm.device)

        outputs = self.llm.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.3,
            top_p=0.9
        )

        response = self.llm_tokenizer.decode(
            outputs[0][len(inputs.input_ids[0]):],
            skip_special_tokens=True
        )

        return response.strip()

